import React, { useState, useEffect } from 'react';
import { Layout, Menu, Divider, Button } from 'antd';
import {
  UploadOutlined,
  DatabaseOutlined,
  FileSearchOutlined,
  SettingOutlined,
  CompressOutlined,
  PlayCircleOutlined,
  ThunderboltOutlined,
  RedoOutlined
} from '@ant-design/icons';
import ModelManage from './ModelManage01';
import UploadPage from './UploadPage';
import ModelCompression from './ModelCompression';
import ModelTraining from './ModelTraining'; // 新增：导入模型训练组件
import { useNavigate, useLocation } from 'react-router-dom';
import AlgorithmMange from './AlgorithmMange';
import DataManage from './DataManage';
import ModelUpdate from './ModelUpdate';

const { Header, Content, Sider } = Layout;

const Management = () => {
  const [selectedKey, setSelectedKey] = useState('modelManage');
  const navigate = useNavigate();
  const location = useLocation();

  useEffect(() => {
    if (location.state && location.state.selectedKey) {
      setSelectedKey(location.state.selectedKey);
    }
  }, [location.state]);

  const handleBackClick = () => {
    navigate('/');
    window.location.reload();
  };

  const handleMenuClick = (e) => {
    setSelectedKey(e.key);
  };

  const renderContent = () => {
    switch (selectedKey) {
      case 'modelManage':
        return <ModelManage />;
      case 'uploadPage':
        return <UploadPage />;
      case 'modelTraining':  // 新增：添加模型训练case
        return <ModelTraining />;
      case 'modelCompression':
        return <ModelCompression />;
      case 'compressionLibrary':
        return <AlgorithmMange />;
      case 'dataManagement':
        return <DataManage />;
      case 'modelUpdate':
        return <ModelUpdate />
      default:
        return <ModelManage />;
    }
  };

  return (
    <Layout style={{ height: '95vh', marginTop: '0', backgroundColor: 'rgba(0, 0, 0, 0.8)' }}>
      <Header style={{ color: 'white', fontSize: '24px', fontFamily: 'Orbitron', border: '1px solid rgba(126, 156, 179, 0.8)' }}>
        天算星座-模型上星管理平台
        <Button type='primary' onClick={handleBackClick} style={{ float: 'right', marginTop: '16px' }}>
          返回
        </Button>
      </Header>
      <Layout style={{ backgroundColor: 'transparent' }}>
        <Sider width={200} className="site-layout-background" style={{ backgroundColor: 'rgba(24, 73, 110, 0.47)', border: '1px solid rgba(126, 156, 179, 0.8)' }}>
          <Menu
            mode="inline"
            selectedKeys={[selectedKey]}
            style={{ height: '100%', borderRight: 0 }}
            onClick={handleMenuClick}
            theme='dark'
          >
            <Menu.Item key="modelManage" icon={<FileSearchOutlined />}>
              模型管理
            </Menu.Item>
            {/*<Menu.Item key="uploadPage" icon={<UploadOutlined />}>*/}
            {/*  模型上传*/}
            {/*</Menu.Item>*/}
            <Menu.Item key="modelTraining" icon={<ThunderboltOutlined />}>
              模型训练
            </Menu.Item>
            <Menu.Item key="modelCompression" icon={<CompressOutlined />}>
              模型压缩
            </Menu.Item>
            {/*<Menu.Item key="compressionLibrary" icon={<DatabaseOutlined />}>*/}
            {/*  压缩算法库*/}
            {/*</Menu.Item>*/}
            {/*<Menu.Item key="dataManagement" icon={<SettingOutlined />}>*/}
            {/*  数据集管理*/}
            {/*</Menu.Item>*/}
            <Menu.Item key="modelInference" icon={<PlayCircleOutlined />}>
              模型推理
            </Menu.Item>
            <Menu.Item key="modelUpdate" icon={<RedoOutlined />}>
              模型更新
            </Menu.Item>
          </Menu>
        </Sider>
        <Layout style={{ padding: '0', backgroundColor: 'rgba(24, 73, 110, 0.47)', border: '1px solid rgba(126, 156, 179, 0.8)' }}>
          <Content
            style={{
              padding: 24,
              margin: 0,
              minHeight: 280,
              backgroundColor: 'rgba(255,255,255,0.27)',
            }}
          >
            {renderContent()}
          </Content>
        </Layout>
      </Layout>
    </Layout>
  );
};

export default Management;